/**
 * Copyright (c) 2011, Tim Buktu
 * All rights reserved.
 * <p>
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 * <p>
 * 1. Redistributions of source code must retain the above copyright notice,
 * this list of conditions and the following disclaimer.
 * 2. Redistributions in binary form must reproduce the above copyright
 * notice, this list of conditions and the following disclaimer in the
 * documentation and/or other materials provided with the distribution.
 * <p>
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
 * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
 * ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
 * LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
 * CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
 * SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
 * INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
 * CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
 * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
 * POSSIBILITY OF SUCH DAMAGE.
 */

package com.example.emailWeb.ntru.encrypt;

import com.example.emailWeb.ntru.encrypt.EncryptionParameters.TernaryPolynomialType;
import com.example.emailWeb.ntru.exception.NtruException;
import com.example.emailWeb.ntru.polynomial.*;

import java.nio.ByteBuffer;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.security.SecureRandom;
import java.util.Arrays;
import java.util.Random;

/**
 * Encrypts, decrypts data and generates key pairs.<br/>
 * The parameter p is hardcoded to 3.
 */
public class myNtruEncrypt {
    private EncryptionParameters params;

    /**
     * Constructs a new instance with a set of encryption parameters.
     *
     * @param params encryption parameters
     */
    public myNtruEncrypt(EncryptionParameters params) {
        this.params = params;
    }

    /**
     * Generates a new encryption key pair using two threads if possible.
     *
     * @return a key pair
     */
    public EncryptionKeyPair generateKeyPair() {
        return generateKeyPair(new SecureRandom(), true);
    }

    /**
     * Generates a new encryption key pair in a single thread.
     *
     * @return a key pair
     */
    public EncryptionKeyPair generateKeyPairSingleThread() {
        return generateKeyPair(new SecureRandom(), false);
    }

    /**
     * Generates an encryption key pair from a passphrase using two threads if possible.<br/>
     * Invoking this method with the same passphrase and salt will always return the
     * same key pair.
     *
     * @param passphrase
     * @param salt       salt for the passphrase; can be <code>null</code> but this is strongly discouraged
     * @return a key pair
     */
    public EncryptionKeyPair generateKeyPair(char[] passphrase, byte[] salt) {
        PassphraseBasedPRNG rng = new PassphraseBasedPRNG(passphrase, salt);
        return generateKeyPair(rng, rng.createBranch(), true);
    }

    /**
     * Generates an encryption key pair from a passphrase in a single thread.<br/>
     * Invoking this method with the same passphrase and salt will always return the
     * same key pair.
     *
     * @param passphrase
     * @param salt       salt for the passphrase; can be <code>null</code> but this is strongly discouraged
     * @return a key pair
     */
    public EncryptionKeyPair generateKeyPairSingleThread(char[] passphrase, byte[] salt) {
        Random rng = new PassphraseBasedPRNG(passphrase, salt);
        return generateKeyPair(rng, false);
    }

    /**
     * A convenience method that generates a random 128-bit salt vector for key pair generation.
     *
     * @return a new salt vector
     */
    public byte[] generateSalt() {
        byte[] salt = new byte[16];
        new SecureRandom().nextBytes(salt);
        return salt;
    }

    private EncryptionKeyPair generateKeyPair(Random rng, boolean multiThread) {
        return generateKeyPair(rng, rng, multiThread);
    }

    /**
     * Generates a new encryption key pair.
     *
     * @param rngf        the random number generator to use for generating the secret polynomial f
     * @param rngg        the random number generator to use for generating the secret polynomial g
     * @param multiThread whether to use two threads; only has an effect if more than one virtual processor is available
     * @return a key pair
     */
    private EncryptionKeyPair generateKeyPair(Random rngf, final Random rngg, boolean multiThread) {
        int N = params.N;
        int q = params.q;
        int df = params.df;
        int df1 = params.df1;
        int df2 = params.df2;
        int df3 = params.df3;
        boolean fastFp = params.fastFp;
        boolean sparse = params.sparse; // 疏
        TernaryPolynomialType polyType = params.polyType; // SIMPlE or PRODUCT

        Polynomial t;
        IntegerPolynomial fq;
        IntegerPolynomial fp = null;

        // Choose a random g that is invertible mod q. Start a new thread if multiThread=true and more than one processor is available.
        // g (mod q) 有逆
        //选择一个可逆模 q 的随机 g。 如果 multiThread=true 并且有多个处理器可用，则启动一个新线程

        IntegerPolynomial g = null;
        g = generateG(rngg);

        // choose a random f that is invertible mod 3 and q
        IntegerPolynomial f;
        while (true) {


            // choose random t, calculate f and fp

            t = polyType == TernaryPolynomialType.SIMPLE ?
                    PolynomialGenerator.generateRandomTernary(N, df, df - 1, sparse, rngf) :
                    ProductFormPolynomial.generateRandom(N, df1, df2, df3, df3 - 1, rngf);
            f = t.toIntegerPolynomial();
            fp = f.invertF3(); //当fastFp为false时，f * fp(mod 3) = 1;
            //问：fp 模3, 而fq模q?
            if (fp == null)
                continue;
            System.out.println("fp: " + Arrays.toString(fp.coeffs));

            fq = f.invertFq(q); //q必须是2的倍数；f * fq (mod q) = 1;
            if (fq != null)
                break;
        }
        System.out.println("f: " + t.toIntegerPolynomial());
        System.out.println("fq: " + fq);
        System.out.println("fp: " + fp);

        System.out.println("f * fq(mod q):");
        IntegerPolynomial Ffq = f.mult(fq,q);
        System.out.println(Ffq);

        IntegerPolynomial h = g.mult(fq, q); //h是公钥; h = g*fq(mod q)
        h.mult3(q);         //问：h = h*3 (mod q) ? 答：p = 3; 故 h = p * g * fq(mod q);
//        h.mult(3);

        System.out.println("h before positive: ");
        System.out.println(h);

        h.ensurePositive(q);

        System.out.println("h after positive: ");
        System.out.println(h);
        IntegerPolynomial hClone = h.clone();
        hClone.mod3();
        System.out.println("h mod 3: " + hClone);


        IntegerPolynomial hNotModQ = g.mult(fq);
        hNotModQ.mult(3);
        System.out.println("h not mod q: ");
        System.out.println(hNotModQ);

        g.clear(); //coeffs[]元素全部设置为0；
        fq.clear();

        EncryptionPrivateKey priv = new EncryptionPrivateKey(t, fp, N, q, sparse, fastFp, polyType); // t是f的三元多项式
        EncryptionPublicKey pub = new EncryptionPublicKey(h, N, q);
        return new EncryptionKeyPair(priv, pub);
    }

    /**
     * Generates the ephemeral(临时) secret polynomial <code>g</code>.
     * g模2可逆
     *
     * @return
     */
    private IntegerPolynomial generateG(Random rng) {
        final int N = params.N;
        int dg = params.dg; // dg = N / 3;

        while (true) {
            DenseTernaryPolynomial g = DenseTernaryPolynomial.generateRandom(N, dg, dg - 1, rng);
            //产生N个数，dg个1，dg-1个-1，其余的为0，rng用作shuffle的参数。
            if (g.isInvertiblePow2())
                return g;
        }
    }

    /**
     * Encrypts a message.<br/>
     * See P1363.1 section 9.2.2.
     *
     * @param m      The message to encrypt
     * @param pubKey the public key to encrypt the message with
     * @return the encrypted message
     * @throws NtruException if the JRE doesn't implement the specified hash algorithm, the message is longer than <code>maxLenBytes</code>, or <code>maxLenBytes</code> is greater than 255
     */
    public byte[] encrypt(byte[] m, EncryptionPublicKey pubKey) {
        IntegerPolynomial pub = pubKey.h;
        int N = params.N;
        int q = params.q;
        int maxLenBytes = params.maxMsgLenBytes;
        int db = params.db; // 不同的加密模式，有128 bit和256 bit。
        int bufferLenBits = params.bufferLenBits;
        int dm0 = params.dm0;
        int maxM1 = params.maxM1;
        int minCallsMask = params.minCallsMask;
        boolean hashSeed = params.hashSeed;

        int l = m.length;
        if (maxLenBytes > 255)
            throw new NtruException("llen values bigger than 1 are not supported");
        if (l > maxLenBytes)
            throw new NtruException("Message too long: " + l + ">" + maxLenBytes);

        SecureRandom rng = new SecureRandom();
        while (true) {
            // M = b|octL|m|p0
            byte[] b = new byte[db / 8]; //db: number of random bits to prepend to the message; should be a multiple of 8
            rng.nextBytes(b);
            byte[] p0 = new byte[maxLenBytes + 1 - l];
            ByteBuffer mBuf = ByteBuffer.allocate((bufferLenBits + 7) / 8);
            mBuf.put(b);
            mBuf.put((byte) l); // l是明文消息的字节长度
            mBuf.put(m); //
            mBuf.put(p0);
            byte[] M = mBuf.array();
            //SimpleExample.java: maxM1 = 126,mTrin由 -1,0,1组成的int[]
            IntegerPolynomial mTrin = IntegerPolynomial.fromBinary3Sves(M, N, maxM1 > 0);   // don't use the constant coeff if maxM1 is set; see below
//            IntegerPolynomial mTrin = IntegerPolynomial.fromBinary3Sves(m, N, maxM1>0);   // don't use the constant coeff if maxM1 is set; see below

            byte[] sData = getSeed(m, pub, b);
            //从多项式环 Ř 中随机选取小多项式 r.
            Polynomial r = generateBlindingPoly(sData);

            System.out.println("r(小多项式):" + r.toIntegerPolynomial());

            IntegerPolynomial R = r.mult(pub, q); // 把pub看作h, R = r*h(mod q)
            mTrin.mod3();
            // 原来：R = h * r(mod q)；现在 R+m(mod q);
            // 故：R = h*r+m(mod q); 把R用C代替，即 C = h*r+m(mod q);
            R.add(mTrin, q);
            // APR2011_439_FAST : q = 2048; 2^11 = 2048;
            // ensurePositive正数化的原因：在转换成byte的过程中，使用11个bit即可表示，可以省下大量的空间。
            System.out.println("R(C密文) before positive:");
            System.out.println(R);
            R.ensurePositive(q);
            System.out.println("R(C密文) after positive:");
            System.out.println(R);

            System.out.println("R(C密文) not mod q: ");
            IntegerPolynomial C = r.mult(pub);
            C.add(mTrin);
            System.out.println(C);

            return R.toBinary(q);
        }
    }

    /**
     * Generates a seed for the Blinding Polynomial Generation Function.
     *
     * @param m   the plain-text message
     * @param pub the public key
     * @param b   <code>db</code> bits of random data
     * @return a byte array containing a seed value
     */
    private byte[] getSeed(byte[] m, IntegerPolynomial pub, byte[] b) {
        byte[] oid = params.oid;

        byte[] hTrunc = pub.toBinaryTrunc(params.q, params.pkLen / 8);
        // sData = OID|m|b|hTrunc
        byte[] sData = new byte[oid.length + m.length + b.length + hTrunc.length];
        System.arraycopy(oid, 0, sData, 0, oid.length);
        int start = oid.length;
        System.arraycopy(m, 0, sData, start, m.length);
        start += m.length;
        System.arraycopy(b, 0, sData, start, b.length);
        start += b.length;
        System.arraycopy(hTrunc, 0, sData, start, hTrunc.length);
        return sData;
    }

    /**
     * Deterministically generates a blinding polynomial from a seed and a message representative.
     *
     * @param seed
     * @return a blinding polynomial
     */
    private Polynomial generateBlindingPoly(byte[] seed) {
        int N = params.N;
        IndexGenerator ig = new IndexGenerator(seed, params);

        if (params.polyType == TernaryPolynomialType.PRODUCT) {
            SparseTernaryPolynomial r1 = SparseTernaryPolynomial.generateBlindingPoly(ig, N, params.dr1);
            SparseTernaryPolynomial r2 = SparseTernaryPolynomial.generateBlindingPoly(ig, N, params.dr2);
            SparseTernaryPolynomial r3 = SparseTernaryPolynomial.generateBlindingPoly(ig, N, params.dr3);
            return new ProductFormPolynomial(r1, r2, r3);
        } else if (params.sparse)
            return SparseTernaryPolynomial.generateBlindingPoly(ig, N, params.dr);
        else
            return DenseTernaryPolynomial.generateBlindingPoly(ig, N, params.dr);
    }

    /**
     * An implementation of MGF-TP-1 from P1363.1 section 8.4.1.1.
     *
     * @param seed
     * @param N
     * @param minCallsMask
     * @param hashSeed     whether to hash the seed
     * @return
     * @throws NtruException if the JRE doesn't implement the specified hash algorithm
     */
    private IntegerPolynomial MGF(byte[] seed, int N, int minCallsMask, boolean hashSeed) {
        MessageDigest hashAlg;
        try {
            hashAlg = MessageDigest.getInstance(params.hashAlg);
        } catch (NoSuchAlgorithmException e) {
            throw new NtruException(e);
        }

        int hashLen = hashAlg.getDigestLength();
        ByteBuffer buf = ByteBuffer.allocate(minCallsMask * hashLen);
        byte[] Z = hashSeed ? hashAlg.digest(seed) : seed;
        int counter = 0;
        while (counter < minCallsMask) {
            ByteBuffer hashInput = ByteBuffer.allocate(Z.length + 4);
            hashInput.put(Z);
            hashInput.putInt(counter);
            byte[] hash = hashAlg.digest(hashInput.array());
            buf.put(hash);
            counter++;
        }

        IntegerPolynomial i = new IntegerPolynomial(N);
        while (true) {
            int cur = 0;
            for (byte o : buf.array()) {
                int O = (int) o & 0xFF;
                if (O >= 243)   // 243 = 3^5
                    continue;

                for (int terIdx = 0; terIdx < 4; terIdx++) {
                    int rem3 = O % 3;
                    i.coeffs[cur] = rem3 == 2 ? -1 : rem3;   /* reduce to [-1..1] */
                    cur++;
                    if (cur == N)
                        return i;
                    O = (O - rem3) / 3;
                }

                i.coeffs[cur] = O == 2 ? -1 : O;   /* reduce to [-1..1] */
                cur++;
                if (cur == N)
                    return i;
            }

            if (cur >= N)
                return i;

            buf = ByteBuffer.allocate(hashLen);
            ByteBuffer hashInput = ByteBuffer.allocate(Z.length + 4);
            hashInput.put(Z);
            hashInput.putInt(counter);
            byte[] hash = hashAlg.digest(hashInput.array());
            buf.put(hash);
            counter++;
        }
    }

    /**
     * Decrypts a message.<br/>
     * See P1363.1 section 9.2.3.
     *
     * @param data The message to decrypt
     * @param kp   a key pair that contains the public key the message was encrypted with, and the corresponding private key
     * @return the decrypted message
     * @throws NtruException if the JRE doesn't implement the specified hash algorithm, the encrypted data is invalid, or <code>maxLenBytes</code> is greater than 255
     */
    public byte[] decrypt(byte[] data, EncryptionKeyPair kp) {
        Polynomial priv_t = kp.priv.t; // t 就是 f，私钥

        // 若 fastFp = true,则fp = 1;
        // 若 fastFp = false,则 f * fp(mod p) = 1;
        IntegerPolynomial priv_fp = kp.priv.fp; // f * fp(mod p) = 1; 可用fp和f相乘削去f。
        IntegerPolynomial pub = kp.pub.h;
        int N = params.N;
        int q = params.q;
        int db = params.db;
        int maxMsgLenBytes = params.maxMsgLenBytes;
        int dm0 = params.dm0;
        int maxM1 = params.maxM1;
        int minCallsMask = params.minCallsMask;
        boolean hashSeed = params.hashSeed;

        if (maxMsgLenBytes > 255)
            throw new NtruException("maxMsgLenBytes values bigger than 255 are not supported");

        int bLen = db / 8; // number of random bits to prepend to the message; should be a multiple of 8

        // e : R + mTrin (mod q)
        IntegerPolynomial e = IntegerPolynomial.fromBinary(data, N, q);//data从字节变成int[]

        System.out.println("从密文字节数组解码出的多项式: ");
        System.out.println(e);

        IntegerPolynomial ci = decrypt(e, priv_t, priv_fp); // mTrin

        IntegerPolynomial cMTrin = ci;
        cMTrin.mod3(); // 为什么一直要mod3?
        byte[] cM = cMTrin.toBinary3Sves(maxM1 > 0);

        ByteBuffer buf = ByteBuffer.wrap(cM);
        byte[] cb = new byte[bLen]; // bLen = db / 8;
        buf.get(cb);
        int cl = buf.get() & 0xFF;   // llen=1, so read one byte; 代表明文的消息的字节长度；
        if (cl > maxMsgLenBytes)
            throw new NtruException("Message too long: " + cl + ">" + maxMsgLenBytes);
        byte[] cm = new byte[cl];
        buf.get(cm); // 得到明文

        System.out.println("cm in decrypy:");
        System.out.println(new String(cm));

//        byte[] p0 = new byte[buf.remaining()];
//        buf.get(p0);
//        if (!Arrays.equals(p0, new byte[p0.length]))
//            throw new NtruException("The message is not followed by zeroes");
//
//        byte[] sData = getSeed(cm, pub, cb);
//
//        Polynomial cr = generateBlindingPoly(sData);
//        IntegerPolynomial cRPrime = cr.mult(pub);
//        cRPrime.modPositive(q);
//        if (!cRPrime.equals(cR))
//            throw new NtruException("Invalid message encoding");
        return cm;
    }

    /**
     * @param e
     * @param priv_t  a polynomial such that if <code>fastFp=true</code>, <code>f=1+3*priv_t</code>; otherwise, <code>f=priv_t</code>
     * @param priv_fp
     * @return
     */
    IntegerPolynomial decrypt(IntegerPolynomial e, Polynomial priv_t, IntegerPolynomial priv_fp) {
        // e 看作 C, priv_t 看作 f , priv_fp 看作 fp；
        // C = h*r + m(mod q);
        int q = params.q;

        IntegerPolynomial a;

            // f * C (mod q) = f*h*r + f*m(mod q);
            // h = p*g*fq(mod q); f*fq(mod q) = 1;  --> f*h(mod q) = f*p*g*fq(mod q) = p*g;
            //故：f * C (mod q) = p*g*r + f*m(mod q);

        a = priv_t.mult(e, q); // e*f (mod q)

        System.out.println("f*C(mod q): ");
        System.out.println(a);

        System.out.println("a not center0(q) mod 3:");
        IntegerPolynomial tmp = a.clone();
        tmp.mod3();
        System.out.println(tmp);
        a.center0(q);
        System.out.println("a.center0(q): "+a);

        a.mod3(); // f*m
        System.out.println("f*m(mod q)(mod p): "+a);
        System.out.println("DenseTernaryPolynomial(a): "+new DenseTernaryPolynomial(a));
        IntegerPolynomial c = new DenseTernaryPolynomial(a).mult(priv_fp, 3);
        c.center0(3);
//        c.mod3();
        return c;
    }
}